DataFrame Column Modification Operations

Manipulating DataFrame schemas in PySpark by adding, casting, renaming, or dropping columns.

What are the Schema Modification Operations?

A DataFrame is structurally immutable. In PySpark, we use structural DSL methods to generate a new DataFrame with modified schemas:

• withColumn(colName, colExpression): Appends a new column or replaces an existing one if the name matches.
• withColumnRenamed(oldName, newName): Alters the name of a specific column.
• drop(colNames...): Excludes specific columns from the DataFrame projection.
• cast(dataType): Changes the data type of an existing column.

Syntax and Common Scenarios

from pyspark.sql.functions import col
from pyspark.sql.types import IntegerType

# A. Add a column
df.withColumn("is_active", col("status") == "ACTIVE")

# B. Cast data type
df.withColumn("age", col("age").cast(IntegerType()))

# C. Rename column
df.withColumnRenamed("first_name", "first")

# D. Drop multiple columns
df.drop("address", "zip_code")

Example Usage Pipeline

Below is a complete, copy-paste-ready PySpark script demonstrating schema manipulations:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when
from pyspark.sql.types import IntegerType

# 1. Setup local Spark session
spark = SparkSession.builder \
    .appName("DataFrame Schema Manipulations") \
    .master("local[*]") \
    .getOrCreate()

# 2. Dummy dataset with unclean inputs
data = [
    ("1001", "Alice", "92000", "ACTIVE"),
    ("1002", "Bob", "61000", "SUSPENDED"),
    ("1003", "Charlie", "48000", "ACTIVE"),
]
columns = ["id", "name", "raw_salary", "status"]
df = spark.createDataFrame(data, columns)

# 3. Clean and transform columns:
# - Rename 'name' to 'employee name'
# - Cast 'raw salary' to Integer and calculate net salary (bonus/tax)
# - Add conditional column 'eligible for bonus' based on active status
# - Drop the legacy 'raw salary' and 'status' columns
transformed_df = df \
    .withColumnRenamed("name", "employee_name") \
    .withColumn("salary", col("raw_salary").cast(IntegerType())) \
    .withColumn("net_income", col("salary") * 0.90) \
    .withColumn("bonus_eligible", when(col("status") == "ACTIVE", True).otherwise(False)) \
    .drop("raw_salary", "status")

# 4. Show results
print("=== Original Schema and Data ===")
df.show()
df.printSchema()

print("=== Transformed Schema and Data ===")
transformed_df.show()
transformed_df.printSchema()

Rendered Output:

=== Original Schema and Data ===
+----+-------+----------+---------+
|  id|   name|raw_salary|   status|
+----+-------+----------+---------+
|1001|  Alice|     92000|   ACTIVE|
|1002|    Bob|     61000|SUSPENDED|
|1003|Charlie|     48000|   ACTIVE|
+----+-------+----------+---------+

root
 |-- id: string (nullable = true)
 |-- name: string (nullable = true)
 |-- raw_salary: string (nullable = true)
 |-- status: string (nullable = true)

=== Transformed Schema and Data ===
+----+-------------+------+----------+--------------+
|  id|employee_name|salary|net_income|bonus_eligible|
+----+-------------+------+----------+--------------+
|1001|        Alice| 92000|   82800.0|          true|
|1002|          Bob| 61000|   54900.0|         false|
|1003|      Charlie| 48000|   43200.0|          true|
+----+-------------+------+----------+--------------+

root
 |-- id: string (nullable = true)
 |-- employee_name: string (nullable = true)
 |-- salary: integer (nullable = true)
 |-- net_income: double (nullable = true)
 |-- bonus_eligible: boolean (nullable = false)